#include "TrainAlgs.h"

#include "Algorithms.h"

#include <stdlib.h>
#include <stdio.h>
#include <iostream>
#include <fstream>
#include <assert.h>
#include <cmath>

void HMMALG::trainModel(const Observation &obs, unsigned n, unsigned m,
						std::vector<double> &pi,
						std::vector<std::vector<double> > &a,
						std::vector<std::vector<double> > &b,
						unsigned max_iter, double delta, double eps) {

	size_t tmax = obs.size();
	
	std::vector<double> normfactors(tmax);
	std::vector<std::vector<double> > scaled_alpha(tmax);
	std::vector<std::vector<double> > scaled_beta(tmax);
	std::vector<std::vector<double> > gamma(tmax);
	
	for (size_t i = 0; i < tmax ; i++) {
		scaled_alpha[i].resize(n);
		scaled_beta[i].resize(n);
		gamma[i].resize(n);
	}
	/*
	printf("pi\n");
	printVector(pi);
	printf("a\n");
	printMatrix(a);
	printf("b\n");
	printMatrix(b);
	*/
	double logProb0 = -1.0e+10;
	double logProb1;
	
	std::vector<std::vector<double> > a_save;
	std::vector<std::vector<double> > b_save;
	std::vector<double> pi_save;
	
	// debug
	//std::vector<unsigned> path;
	
	for(unsigned iter = 0; iter < max_iter; iter++) {
		
		calculateScaledAlpha(obs, n,  m, pi, a, b, scaled_alpha, normfactors);

		logProb1 = logProb(normfactors);
		
		printf("Iter = %d\tLOGPROB = %f\n", iter, logProb1);
		if (logProb1 - logProb0 < delta || isnan(logProb1)) {
			a = a_save;
			b = b_save;
			pi = pi_save;
			break;
		}
		/*
		findMostProbablePath(obs, n, m, pi, a, b, path);
		printf("path:\n");
		for (size_t t = 0; t < path.size(); t++) {
			printf("%d, ", (int)path[t]);
		}
		printf("\n");
		// }
		*/
		
		a_save = a;
		b_save = b;
		pi_save = pi;
		
		calculateScaledBeta(obs, n,  m, pi, a, b, scaled_beta);
		calculateGamma(scaled_alpha, scaled_beta, gamma);

/*
		if (iter >= 20 && iter <= 21) {
			printf("alpha:\n");
			printMatrix(scaled_alpha);
			printf("beta:\n");
			printMatrix(scaled_beta);
			printf("gamma:\n");
			printMatrix(gamma);
		}
*/

	   // calculate ~pi
		for (size_t i = 0; i < n; i++) {
			pi[i] = gamma[0][i];
		}
		// calculate ~a
		for (size_t i = 0; i < n; i++) {
			for (size_t j = 0; j < n; j++) {
				double numerator = 0;
				double denominator = 0;
				for (size_t t = 0; t < tmax - 1; t++) {
					numerator += scaled_alpha[t][i] * a[i][j] * scaled_beta[t + 1][j] * b[j][obs[t + 1]];
					denominator += scaled_alpha[t][i] * scaled_beta[t][j];
				}
				a[i][j] = numerator / denominator;
				//a[i][j] = (numerator < eps || denominator < eps) ? 0 : numerator / denominator;
			}
			normalize(a[i], eps);
		}

		// calculate ~b
		for (size_t i = 0; i < n; i++) {
			for (size_t k = 0; k < m; k++) {
				double numerator = 0;
				double denominator = 0;
				for (size_t t = 0; t < tmax; t++) {
					if (obs[t] == k) {
						numerator += gamma[t][i];
					}
					denominator += gamma[t][i];
				}
				b[i][k] = numerator / denominator;
			}
			normalize(b[i], eps);
		}
		logProb0 = logProb1;
/*		
		if (iter >= 20 && iter <= 21) 
			printf("iteration # %d \n", iter);
			printf("pi\n");
			printVector(pi);
			printf("a\n");
			printMatrix(a);
			printf("b\n");
			printMatrix(b);
		}
*/
	}
/*
			printf("pi\n");
			printVector(pi);
			printf("a\n");
			printMatrix(a);
	printf("traning finished!\n");
*/			
}



void HMMALG::trainModel(const std::vector<Observation> &obs, unsigned n, unsigned m,
						std::vector<double> &pi,
						std::vector<std::vector<double> > &a,
						std::vector<std::vector<double> > &b,
						unsigned max_iter, double delta, double eps) {

	unsigned  nobs = obs.size();
	
	// prepare matrices
	std::vector<double> normfactors;
	std::vector<std::vector<std::vector<double> > > scaled_alpha(nobs);
	std::vector<std::vector<std::vector<double> > > scaled_beta(nobs);
	std::vector<std::vector<std::vector<double> > > gamma(nobs);
	std::vector<double> logProbs(nobs);
	
	for (unsigned i = 0; i < nobs ; i++) {
		unsigned tmax = obs[i].size();
		scaled_alpha[i].resize(tmax);
		scaled_beta[i].resize(tmax);
		gamma[i].resize(tmax);
		for (unsigned j = 0; j < tmax; j++) {
			scaled_alpha[i][j].resize(n);
			scaled_beta[i][j].resize(n);
			gamma[i][j].resize(n);
		}
	}
	
	// difine variables
	double logProb0 = -1.0e+10;
	double logProb1;
	
	int cnt = 0;
	
	std::vector<std::vector<double> > a_save;
	std::vector<std::vector<double> > b_save;
	std::vector<double> pi_save;
	
	for(unsigned iter = 0; iter < max_iter; iter++) {
	
		for (unsigned i = 0; i < nobs; i++) {
			normfactors.resize(obs[i].size());
			calculateScaledAlpha(obs[i], n,  m, pi, a, b, scaled_alpha[i], normfactors);
			logProbs[i] = logProb(normfactors);
		}
		
		logProb1 = sumVector(logProbs);
		
		std::cout << "ITER: " << iter<<" logProbs:\n";
		for (unsigned i = 0; i < nobs; i++) {
			std::cout << logProbs[i] << "\n";
		}
		std::cout << "TOTAL logProb: " << logProb1 << "\n";
		
		/*
		printf("pi\n");
		printVector(pi);
		printf("a\n");
		printMatrix(a);
		printf("b\n");
		printMatrix(b);
		*/
		if (logProb1 - logProb0 < delta || isnan(logProb1)) {
			cnt++;
			if (cnt >= 2) {
				a = a_save;
				b = b_save;
				pi = pi_save;
				break;
			}
		} else {
			a_save = a;
			b_save = b;
			pi_save = pi;
			cnt = 0;
		}
		
		
		for (unsigned i = 0; i < nobs; i++) {
			calculateScaledBeta(obs[i], n,  m, pi, a, b, scaled_beta[i]);
			calculateGamma(scaled_alpha[i], scaled_beta[i], gamma[i]);
		}

		//calculate ~pi
		for (unsigned  i = 0; i < n; i++) {
			for (unsigned k = 0; k < nobs; k++) {
				pi[i] = gamma[k][0][i] * exp(-logProbs[k]);
			}
		}
		
		normalize(pi, eps);
		
		// calculate ~a
		for (unsigned  i = 0; i < n; i++) {
			for (unsigned  j = 0; j < n; j++) {
				double numerator = 0;
				double denominator = 0;
				for (unsigned  k = 0; k < nobs; k++) {
					unsigned  tmax = obs[k].size();
					for (unsigned  t = 0; t < tmax - 1; t++) {
						numerator += exp(-logProbs[k]) * scaled_alpha[k][t][i] * a[i][j] * scaled_beta[k][t + 1][j] * b[j][obs[k][t + 1]];
						denominator += exp(-logProbs[k]) * scaled_alpha[k][t][i] * scaled_beta[k][t][j];
					}
				}
				a[i][j] = numerator / denominator;
				//a[i][j] = (numerator < eps || denominator < eps) ? 0 : numerator / denominator;
			}
			normalize(a[i], eps);
		}

		// calculate ~b
		for (unsigned  i = 0; i < n; i++) {
			for (unsigned  j = 0; j < m; j++) {
				double numerator = 0;
				double denominator = 0;
				for (unsigned  k = 0; k < nobs; k++) {
					unsigned  tmax = obs[k].size();
					for (unsigned  t = 0; t < tmax; t++) {
						if (obs[k][t] == j) {
							numerator += gamma[k][t][i] * exp(-logProbs[k]);
						}
						denominator += gamma[k][t][i] * exp(-logProbs[k]);
					}
				}
				b[i][j] = numerator / denominator;
			}
			normalize(b[i], eps);
		}
		logProb0 = logProb1;
	}
}

